// LSTM模型结构

// 引入 lodash 库，用于数据处理和操作
const _ = require('lodash')
// 精确计算
const NP = require('number-precision')
// 引入 path 模块，用于处理文件和目录路径
const path = require('path')
// 获取项目根目录
const ROOT_DIR = path.join(__dirname, '..')
// helpers 目录
const HELPERS_DIR = path.join(ROOT_DIR, 'helpers')
// 读取所有数据
let { allValues } = require(path.join(HELPERS_DIR, 'readAllValues.js'))

const tf = require('@tensorflow/tfjs-node')

// 假设这是100期历史开奖数据
const historicalData = allValues.slice(0, 100)

// 数据归一化：将号码映射到 [0, 1] 范围
const normalizeData = (number) => {
  return (number - 1) / 48 // 1-49 范围映射到 0-1
}

// 准备训练数据：使用过去10期数据预测下一期
const prepareData = (historicalData) => {
  const inputData = [] // 存储输入数据
  const outputData = [] // 存储目标输出数据

  // 从第10期开始，每次取过去10期数据作为输入，下一期数据作为输出
  for (let i = 10; i < historicalData.length; i++) {
    // 使用过去10期的号码作为输入
    const input = historicalData.slice(i - 10, i).map((item) => normalizeData(item.number))

    // 下一期的号码作为目标
    const output = normalizeData(historicalData[i].number)

    inputData.push(input) // 将输入数据添加到数组中
    outputData.push(output) // 将输出数据添加到数组中
  }

  // 将 inputData 转换为三维数组: [样本数, 时间步长, 特征数]
  const reshapedInputData = inputData.map((item) => item.map((value) => [value])) // 确保每个时间步是一个数组

  return { inputData: reshapedInputData, outputData } // 返回训练数据
}

// 准备训练数据
const { inputData, outputData } = prepareData(historicalData)

// 将数据转换为 Tensor 格式
const inputTensor = tf.tensor3d(inputData, [inputData.length, 10, 1]) // 三维数据：样本数，时间步长，特征数
const outputTensor = tf.tensor2d(outputData, [outputData.length, 1]) // 二维数据：样本数，输出维度

// 构建 LSTM 模型
const model = tf.sequential()
model.add(
  tf.layers.lstm({
    units: 64,
    activation: 'relu',
    inputShape: [10, 1],
    returnSequences: true,
    dropout: 0.2 // 添加dropout层，防止过拟合
  })
)
model.add(
  tf.layers.lstm({
    units: 32,
    activation: 'relu',
    returnSequences: false,
    dropout: 0.2 // 第二层LSTM也添加dropout
  })
)
model.add(tf.layers.dense({ units: 1, activation: 'linear' })) // 输出层

// 编译模型
model.compile({ optimizer: 'adam', loss: 'meanSquaredError' })

// 训练模型
async function trainModel() {
  console.log('Training data shape: ', inputTensor.shape)
  console.log('Output data shape: ', outputTensor.shape)

  await model.fit(inputTensor, outputTensor, {
    epochs: 100,
    verbose: 0,
    validationSplit: 0.2, // 分出部分数据做验证，避免过拟合
    callbacks: [
      tf.callbacks.earlyStopping({ patience: 10 }) // 移除 restoreBestWeights
    ]
  })
  console.log('模型训练完成')
}

// 预测下一期号码
async function predictNextNumber(inputData) {
  if (inputData.length !== 10) {
    throw new Error('Input data must contain 10 values')
  }

  const inputTensor = tf.tensor3d(inputData, [1, 10, 1]) // 直接传入 inputData
  const prediction = model.predict(inputTensor) // 进行预测
  const predictionArray = prediction.arraySync() // 将预测结果转换为 JavaScript 数组

  // 将归一化的预测结果转回 [1, 49] 范围
  let predictedNumber = Math.round(predictionArray[0][0] * 48) + 1

  // 确保预测值在 [1, 49] 范围内
  predictedNumber = Math.min(49, Math.max(1, predictedNumber))

  return predictedNumber
}

// 主函数
async function main() {
  // 训练模型
  await trainModel()

  // 从历史数据中提取最后10期数据进行预测
  const last10Data = historicalData.slice(-10).map((item) => normalizeData(item.number))

  console.log('Last 10 data: ', last10Data)

  let predictedNumbers = [] // 存储预测的24期号码
  for (let i = 0; i < 24; i++) {
    // 预测下一期的号码
    const nextNumber = await predictNextNumber(last10Data)
    predictedNumbers.push(nextNumber)

    // 将预测的号码添加到数据中，作为下一轮预测的输入
    last10Data.push(normalizeData(nextNumber))
    last10Data.shift() // 保持输入数据的长度为10
  }

  console.log('预测的24期号码：', predictedNumbers)
}

// 调用主函数
main()
